package task

import (
	"fmt"
	"log"
	"net"
	"port-scanner/vars"
	"sync"
	"time"
)

type Task struct {
	IP   string
	Port int
}

func (*Task) Connect(ip string, port int) (net.Conn, error) {
	conn, err := net.DialTimeout("tcp", fmt.Sprintf("%s:%d", ip, port), 5*time.Second)

	defer func() {
		if conn != nil {
			conn.Close()
		}
	}()

	return conn, err
}

func Run(ipList []net.IP, ports []int) {
	taskChan := make(chan Task, vars.GoroutineNum)
	var wg sync.WaitGroup

	for i := 0; i < vars.GoroutineNum; i++ {
		go func() {
			for task := range taskChan {
				_, err := task.Connect(task.IP, task.Port)
				if err == nil {
					log.Printf("ip: %s port: %d is open\n", task.IP, task.Port)
					storeResult(task)
				}
				wg.Done()
			}
		}()
	}

	for _, ip := range ipList {
		for _, port := range ports {
			wg.Add(1)
			taskChan <- Task{ip.String(), port}
		}
	}

	close(taskChan)
	wg.Wait()
}

func storeResult(task Task) {
	if v, ok := vars.Result.Load(task.IP); ok {
		if ports, ok1 := v.([]int); ok1 {
			ports = append(ports, task.Port)
			vars.Result.Store(task.IP, ports)
		}
	} else {
		ports := []int{task.Port}
		vars.Result.Store(task.IP, ports)
	}
}
